import numpy as np
from scipy.spatial.transform import Rotation as R
import quadprog

from ..mathq.useful_math import eul2Rot, Rz3


mpc_N = 10
ch = 3
nx = 12
nu = 13
ncfr_single = 4
ncfr = ncfr_single * 2
ncstxya = 1
ncstxy_single = ncstxya * 4
ncstxy = ncstxy_single * 2
ncstza = 2
ncstz_single = ncstza * 4
ncstz = ncstz_single * 2
nc = ncfr + ncstxy + ncstz

class MPC:
    def __init__(self, dtIn):

        self.m = 80.0
        self.g = -9.8
        self.miu = 0.5
        self.delta_foot = [0.073, 0.125, 0.025, 0.025]

        self.max = [1000.0, 1000.0, -3.0 * self.m * self.g, 1000.0, 1000.0, 1000.0]
        self.min = [-1000.0, -1000.0, 0.0, -1000.0, -1000.0, -1000.0]

        self.Ac = [np.zeros((nu, nu)) for _ in range(mpc_N)]
        self.Bc = [np.zeros((nu, nu)) for _ in range(mpc_N)]
        self.A = [np.zeros((nu, nu)) for _ in range(mpc_N)]
        self.B = [np.zeros((nu, nu)) for _ in range(mpc_N)]
        self.Cc = np.zeros((nc, nu))
        self.C = np.zeros((nc, nu))

        self.Aqp = np.zeros((nu * mpc_N, nu))
        self.Aqp1 = np.zeros((nu * mpc_N, nu * mpc_N))
        self.Bqp1 = np.zeros((nu * mpc_N, nu * mpc_N))
        self.Bqp = np.zeros((nu * mpc_N, nu * ch))
        self.Cqp1 = np.zeros((nu * mpc_N, 1))
        self.Cqp = np.zeros((nu * mpc_N, 1))

        self.Ufe = np.zeros((nu * ch, 1))
        self.Ufe_pre = np.zeros((nu, 1))

        self.Xd = np.zeros((nu * mpc_N, 1))
        self.X_cur = np.zeros((nu, 1))
        self.X_cal = np.zeros((nu, 1))
        self.dX_cal = np.zeros((nu, 1))

        self.L = np.zeros((nu * mpc_N, nu * mpc_N))
        self.K = np.zeros((nu * ch, nu * ch))
        self.M = np.zeros((nu * ch, nu * ch))
        self.alpha = 0.0
        self.H = np.zeros((nu * ch, nu * ch))
        self.c = np.zeros((nu * ch, 1))

        self.u_low = np.zeros((nu * ch, 1))
        self.u_up = np.zeros((nu * ch, 1))
        self.As = np.zeros((nc * ch, nu * ch))
        self.bs = np.zeros((nc * ch, 1))

        self.pCoM = np.zeros((3, 1))
        self.pf2com = np.zeros((6, 1))
        self.pf2comd = np.zeros((6, 1))
        self.pe = np.zeros((6, 1))
        self.pf2comi = np.zeros((mpc_N, 6, 1))
        self.R_curz = np.zeros((mpc_N, 3, 3))
        self.R_cur = np.zeros((3, 3))
        self.Ic = np.zeros((3, 3))

        # self.R_cur = np.eye(3)
        self.R_w2f = np.zeros((3, 3))
        self.R_f2w = np.zeros((3, 3))
        
        self.legState = np.zeros(10, dtype=int)
        
        self.qp_H = np.zeros((nu * ch, nu * ch))
        self.qp_As = np.zeros((nc * ch, nu * ch))
        self.qp_c = np.zeros(nu * ch)
        self.qp_lbA = np.zeros(nc * ch)
        self.qp_ubA = np.zeros(nc * ch)
        self.qp_lu = np.zeros(nu * ch)
        self.qp_uu = np.zeros(nu * ch)
        self.xOpt_iniGuess = np.zeros(nu * ch)

        # self.options = PyOptions()
        # self.options.printLevel = PyPrintLevel.NONE
        self.nWSR = 100
        self.cpu_time = 0.1
        self.qp_Status = 0
        self.qp_nWSR = 0

        self.dt = dtIn
        self.EN = False

        # ncfr_single = 4
        # ncfr = ncfr_single * 2
        # ncstxya = 1
        # ncstxy_single = ncstxya * 4
        # ncstxy = ncstxy_single * 2
        # ncstza = 2
        # ncstz_single = ncstza * 4
        # ncstz = ncstz_single * 2
        # nc = ncfr + ncstxy + ncstz

    def set_weight(self, u_weight, L_diag, K_diag):
        self.alpha = u_weight
        L_diag_N = np.tile(L_diag, mpc_N)
        K_diag_N = np.tile(K_diag, ch)

        self.L = np.diag(L_diag_N)
        self.K = np.diag(K_diag_N)

        for i in range(mpc_N):
            self.R_curz = R.from_euler('z', self.X_cur[2]).as_matrix()
            self.L[i*nu+3:i*nu+6, i*nu+3:i*nu+6] = self.R_curz @ self.L[i*nu+3:i*nu+6, i*nu+3:i*nu+6] @ self.R_curz.T
            self.L[i*nu+6:i*nu+9, i*nu+6:i*nu+9] = self.R_curz @ self.L[i*nu+6:i*nu+9, i*nu+6:i*nu+9] @ self.R_curz.T
            self.L[i*nu+9:i*nu+12, i*nu+9:i*nu+12] = self.R_curz @ self.L[i*nu+9:i*nu+12, i*nu+9:i*nu+12] @ self.R_curz.T

        for i in range(ch):
            self.R_curz = R.from_euler('z', self.X_cur[2]).as_matrix()
            self.K[i*nu:i*nu+3, i*nu:i*nu+3] = self.R_curz @ self.K[i*nu:i*nu+3, i*nu:i*nu+3] @ self.R_curz.T
            self.K[i*nu+3:i*nu+6, i*nu+3:i*nu+6] = self.R_curz @ self.K[i*nu+3:i*nu+6, i*nu+3:i*nu+6] @ self.R_curz.T
            self.K[i*nu+6:i*nu+9, i*nu+6:i*nu+9] = self.R_curz @ self.K[i*nu+6:i*nu+9, i*nu+6:i*nu+9] @ self.R_curz.T
            self.K[i*nu+9:i*nu+12, i*nu+9:i*nu+12] = self.R_curz @ self.K[i*nu+9:i*nu+12, i*nu+9:i*nu+12] @ self.R_curz.T
    
    def construct_constraints(self):
        # 摩擦约束
        Asfr111 = np.array([
            [-1.0, 0.0, -1.0 / np.sqrt(2.0) * self.miu],
            [1.0, 0.0, -1.0 / np.sqrt(2.0) * self.miu],
            [0.0, -1.0, -1.0 / np.sqrt(2.0) * self.miu],
            [0.0, 1.0, -1.0 / np.sqrt(2.0) * self.miu]
        ])
        Asfr11 = Asfr111 @ self.R_w2f
        Asfr1 = np.zeros((ncfr, nu))
        Asfr1[:ncfr_single, :3] = Asfr11
        Asfr1[ncfr_single:, 6:9] = Asfr11
        
        Asfr = np.zeros((ncfr * ch, nu * ch))
        for i in range(ch):
            Asfr[ncfr*i:ncfr*(i+1), nu*i:nu*(i+1)] = Asfr1

        # 力矩约束 x y
        sign_xy = [1.0, -1.0, -1.0, 1.0]
        gxyz = [np.array([0.0, 1.0, 0.0]), np.array([0.0, 1.0, 0.0]),
                np.array([1.0, 0.0, 0.0]), np.array([1.0, 0.0, 0.0])]
        r = [np.array([0.0, 1.0, 0.0]), np.array([0.0, 1.0, 0.0]),
             np.array([1.0, 0.0, 0.0]), np.array([1.0, 0.0, 0.0])]
        p = [np.array([self.delta_foot[0], 0.0, 0.0]),
             np.array([-self.delta_foot[1], 0.0, 0.0]),
             np.array([0.0, self.delta_foot[2], 0.0]),
             np.array([0.0, -self.delta_foot[3], 0.0])]

        Astxy_r = [np.zeros((ncstxya, 6)) for _ in range(4)]
        Astxy11 = np.zeros((ncstxy_single, 6))
        for i in range(4):
            Astxy_r[i][:, :3] = sign_xy[i] * gxyz[i].T @ self.R_w2f @ self.R_f2w @ r[i] @ (self.R_f2w @ r[i]).T @ np.cross(self.R_f2w @ p[i], np.eye(3))
            Astxy_r[i][:, 3:] = sign_xy[i] * gxyz[i].T @ self.R_w2f
            Astxy11[i*ncstxya:(i+1)*ncstxya, :] = Astxy_r[i]

        Astxy1 = np.zeros((ncstxy, nu))
        Astxy1[:ncstxy_single, :6] = Astxy11
        Astxy1[ncstxy_single:, 6:12] = Astxy11

        Astxy = np.zeros((ncstxy * ch, nu * ch))
        for i in range(ch):
            Astxy[ncstxy*i:ncstxy*(i+1), nu*i:nu*(i+1)] = Astxy1

        # 力矩约束 z
        Astz_r = [np.zeros((ncstza, 6)) for _ in range(4)]
        Astz11 = np.zeros((ncstz_single, 6))
        for i in range(4):
            Astz_r[i][:1, :3] = -np.sqrt(np.sum(p[i]**2)) * self.miu * np.array([[0.0, 0.0, 1.0]]) @ self.R_w2f
            Astz_r[i][:1, 3:] = np.array([[0.0, 0.0, 1.0]]) @ self.R_w2f
            Astz_r[i][1:, :3] = Astz_r[i][:1, :3]
            Astz_r[i][1:, 3:] = -Astz_r[i][:1, 3:]
            Astz11[i*ncstza:(i+1)*ncstza, :] = Astz_r[i]

        Astz1 = np.zeros((ncstz, nu))
        Astz1[:ncstz_single, :6] = Astz11
        Astz1[ncstz_single:, 6:12] = Astz11

        Astz = np.zeros((ncstz * ch, nu * ch))
        for i in range(ch):
            Astz[ncstz*i:ncstz*(i+1), nu*i:nu*(i+1)] = Astz1

        # 组合所有约束
        self.As = np.vstack([Asfr, Astxy, Astz])
        self.bs = np.zeros((nc * ch, 1))

    def cal(self):
        if self.EN:
            # QP preparation
            for i in range(mpc_N):
                self.R_curz = R.from_euler('z', self.X_cur[2]).as_matrix()
                self.Ac[i][:3, 6:9] = self.R_curz.T
                self.Ac[i][3:6, 9:12] = np.eye(3)
                self.A[i] = np.eye(nu) + self.dt * self.Ac[i]

            for i in range(mpc_N):
                Ic_W_inv = np.linalg.inv(self.R_curz @ self.Ic @ self.R_curz.T)
                self.Bc[i][6:9, :3] = Ic_W_inv @ np.cross(self.pf2com[:3].T, np.eye(3)).T
                self.Bc[i][6:9, 3:6] = Ic_W_inv
                self.Bc[i][6:9, 6:9] = Ic_W_inv @ np.cross(self.pf2com[3:].T, np.eye(3)).T
                self.Bc[i][6:9, 9:12] = Ic_W_inv
                self.Bc[i][9:12, :3] = np.eye(3) / self.m
                self.Bc[i][9:12, 6:9] = np.eye(3) / self.m
                self.Bc[i][11, 12] = 1.0 / self.m
                self.B[i] = self.dt * self.Bc[i]

            # Construct Aqp and Bqp matrices
            self.Aqp = np.vstack([np.linalg.matrix_power(self.A[0], i+1) for i in range(mpc_N)])

            self.Aqp1 = np.eye(nu * mpc_N)
            for i in range(1, mpc_N):
                for j in range(i):
                    self.Aqp1[i*nu:(i+1)*nu, j*nu:(j+1)*nu] = np.linalg.matrix_power(self.A[0], i-j)

            self.Bqp1 = np.block([[self.B[i] if i==j else np.zeros((nu, nu)) for j in range(mpc_N)] for i in range(mpc_N)])

            Bqp11 = np.zeros((nu * mpc_N, nu * ch))
            Bqp11[:nu*ch, :] = np.eye(nu * ch)
            for i in range(mpc_N - ch):
                Bqp11[nu*ch + i*nu : nu*ch + (i+1)*nu, nu*(ch-1):nu*ch] = np.eye(nu)

            B_tmp = self.Bqp1 @ Bqp11
            self.Bqp = self.Aqp1 @ B_tmp

            # Construct QP matrices
            self.H = 2 * (self.Bqp.T @ self.L @ self.Bqp + self.alpha * self.K)
            self.c = 2 * self.Bqp.T @ self.L @ (self.Aqp @ self.X_cur - self.Xd)

            # Construct constraints
            self.construct_constraints()

            # 使用 quadprog 求解 QP
            G = self.H
            a = -self.c.flatten()
            C = -self.As.T
            b = -self.bs.flatten()

            try:
                solution = quadprog.solve_qp(G, a, C, b, meq=0)
                self.Ufe = solution[0].reshape((nu * ch, 1))
                success = True
            except ValueError as e:
                print(f"QP solving failed: {e}")
                self.Ufe = np.zeros((nu * ch, 1))
                success = False

            if success:
                self.dX_cal = self.Ac[0] @ self.X_cur + self.Bc[0] @ self.Ufe[:nu]
                delta_X = np.zeros((nu, 1))
                delta_X[:3] = 0.5 * self.dX_cal[6:9] * self.dt**2
                delta_X[3:6] = 0.5 * self.dX_cal[9:12] * self.dt**2
                delta_X[6:9] = self.dX_cal[6:9] * self.dt
                delta_X[9:12] = self.dX_cal[9:12] * self.dt

                self.X_cal = (self.Aqp @ self.X_cur + self.Bqp @ self.Ufe)[:nu] + delta_X

            self.Ufe_pre = self.Ufe[:nu]

    def dataBusRead(self, Data):      
        self.X_cur[0:3] = Data.base_rpy.reshape(3, 1)
        self.X_cur[3:6] = Data.q[:3].reshape(3, 1)
        self.X_cur[6:9] = Data.dq[3:6].reshape(3, 1)
        self.X_cur[9:12] = Data.dq[:3].reshape(3, 1)

        if self.EN:
            # Update Xd
            self.Xd[:-nu] = self.Xd[nu:]
            self.Xd[-nu:, 0] = np.concatenate([Data.js_eul_des, Data.js_pos_des, Data.js_omega_des, Data.js_vel_des])
        else:
            Data.js_eul_des = Data.base_rpy
            Data.js_pos_des = Data.base_pos
            Data.js_omega_des = Data.base_omega_W
            Data.js_vel_des = Data.dq[:3]

            for i in range(mpc_N):
                start_idx = i * nx
                self.Xd[start_idx:start_idx+3] = Data.js_eul_des.reshape(3, 1)
                self.Xd[start_idx+3:start_idx+6] = Data.js_pos_des.reshape(3, 1)
                self.Xd[start_idx+6:start_idx+9] = Data.js_omega_des.reshape(3, 1)
                self.Xd[start_idx+9:start_idx+12] = Data.js_vel_des.reshape(3, 1)

        # self.R_cur = eul2Rot(self.X_cur[0], self.X_cur[1], self.X_cur[2])
        self.R_cur = eul2Rot(self.X_cur[0][0], self.X_cur[1][0], self.X_cur[2][0])
        
        for i in range(mpc_N):
            self.R_curz[i] = Rz3(self.X_cur[2])
        self.pCoM = self.X_cur[3:6]
        self.pe[0:3] = Data.fe_l_pos_W.reshape(3, 1)
        self.pe[3:6] = Data.fe_r_pos_W.reshape(3, 1)
        self.pf2com[0:3] = self.pe[0:3] - self.pCoM.reshape(3, 1)
        self.pf2com[3:6] = self.pe[3:6] - self.pCoM.reshape(3, 1)
        self.pf2comd[0:3] = self.pe[0:3] - self.Xd[3:6].reshape(3, 1)
        self.pf2comd[3:6] = self.pe[3:6] - self.Xd[3:6].reshape(3, 1)

        self.Ic = np.array([[12.61, 0, 0.37],
                            [0, 11.15, 0.01],
                            [0.37, 0.01, 2.15]])

        self.legStateCur = Data.legState
        self.legStateNext = Data.legStateNext
        for i in range(mpc_N):
            aa = i * self.dt / 0.4
            phip = Data.phi + aa
            if phip > 1:
                self.legState[i] = self.legStateNext.value
            else:
                self.legState[i] = self.legStateCur.value

        R_slop = R.from_euler('xyz', Data.slop).as_matrix()
        if self.legStateCur == 2:  # Assuming RSt is 2
            self.R_f2w = Data.fe_r_rot_W
        elif self.legStateCur == 1:  # Assuming LSt is 1
            self.R_f2w = Data.fe_l_rot_W
        else:
            self.R_f2w = R_slop
        self.R_w2f = self.R_f2w.T

    def dataBusWrite(self, Data):
        Data.Xd = self.Xd
        Data.X_cur = self.X_cur
        Data.fe_react_tau_cmd = self.Ufe
        Data.X_cal = self.X_cal
        Data.dX_cal = self.dX_cal

        Data.qp_nWSR_MPC = self.nWSR
        Data.qp_cpuTime_MPC = self.cpu_time
        Data.qpStatus_MPC = self.qp_Status

        Data.Fr_ff = self.Ufe[:12, 0]

        k = 5
        Data.des_ddq[:2] = self.dX_cal[9][0], self.dX_cal[11][0]
        Data.des_ddq[5] = k * (self.Xd[6+2] - Data.dq[5])

        Data.des_dq[:3] = self.Xd[9+0][0], self.Xd[9+1][0], self.Xd[9+2][0]
        Data.des_dq[3:5] = [0.0, 0.0] 
        Data.des_dq[5] = self.Xd[6+2][0]

        Data.des_delta_q[:2] = Data.des_dq[:2] * self.dt
        Data.des_delta_q[5] = Data.des_dq[5] * self.dt

        Data.base_rpy_des = np.array([0.005, 0.00, self.Xd[2][0]])
        Data.base_pos_des = self.Xd[3+0][0], self.Xd[3+1][0], self.Xd[3+2][0]
        
    
    def enable(self):
        self.EN = True

    def disable(self):
        self.EN = False

    def get_ENA(self):
        return self.EN